// Copyright 2018-2025 the Deno authors. MIT license.

use std::borrow::Cow;
use std::rc::Rc;

use deno_core::AsyncRefCell;
use deno_core::CancelHandle;
use deno_core::Resource;
use deno_core::ResourceId;
use deno_core::ResourceTable;
use deno_core::error::ResourceError;
use deno_error::JsErrorBox;

use crate::io::TcpStreamResource;
use crate::ops_tls::TlsStreamResource;

pub trait NetworkStreamTrait: Into<NetworkStream> {
  type Resource;
  const RESOURCE_NAME: &'static str;
  fn local_address(&self) -> Result<NetworkStreamAddress, std::io::Error>;
  fn peer_address(&self) -> Result<NetworkStreamAddress, std::io::Error>;
}

#[allow(async_fn_in_trait)]
pub trait NetworkStreamListenerTrait:
  Into<NetworkStreamListener> + Send + Sync
{
  type Stream: NetworkStreamTrait + 'static;
  type Addr: Into<NetworkStreamAddress> + 'static;
  /// Additional data, if needed
  type ResourceData: Default;
  const RESOURCE_NAME: &'static str;
  async fn accept(&self) -> std::io::Result<(Self::Stream, Self::Addr)>;
  fn listen_address(&self) -> Result<Self::Addr, std::io::Error>;
}

/// A strongly-typed network listener resource for something that
/// implements `NetworkListenerTrait`.
pub struct NetworkListenerResource<T: NetworkStreamListenerTrait> {
  pub listener: AsyncRefCell<T>,
  /// Associated data for this resource. Not required.
  #[allow(unused)]
  pub data: T::ResourceData,
  pub cancel: CancelHandle,
}

impl<T: NetworkStreamListenerTrait + 'static> Resource
  for NetworkListenerResource<T>
{
  fn name(&self) -> Cow<'_, str> {
    T::RESOURCE_NAME.into()
  }

  fn close(self: Rc<Self>) {
    self.cancel.cancel();
  }
}

impl<T: NetworkStreamListenerTrait + 'static> NetworkListenerResource<T> {
  pub fn new(t: T) -> Self {
    Self {
      listener: AsyncRefCell::new(t),
      data: Default::default(),
      cancel: Default::default(),
    }
  }

  /// Returns a [`NetworkStreamListener`] from this resource if it is not in use elsewhere.
  fn take(
    resource_table: &mut ResourceTable,
    listener_rid: ResourceId,
  ) -> Result<Option<NetworkStreamListener>, JsErrorBox> {
    if let Ok(resource_rc) = resource_table.take::<Self>(listener_rid) {
      let resource = Rc::try_unwrap(resource_rc)
        .map_err(|_| JsErrorBox::new("Busy", "Listener is currently in use"))?;
      return Ok(Some(resource.listener.into_inner().into()));
    }
    Ok(None)
  }
}

/// Each of the network streams has the exact same pattern for listening, accepting, etc, so
/// we just codegen them all via macro to avoid repeating each one of these N times.
macro_rules! network_stream {
  (
    $(
      $( #[ $meta:meta ] )?
      [$i:ident, $il:ident, $stream:path, $listener:path, $addr:path, $stream_resource:ty, $read_half:ty, $write_half:ty, ]
    ),*
  ) => {
    /// A raw stream of one of the types handled by this extension.
    #[pin_project::pin_project(project = NetworkStreamProject)]
    #[allow(clippy::large_enum_variant)]
    pub enum NetworkStream {
      $(
        $( #[ $meta ] )?
        $i (#[pin] $stream),
      )*
    }

    /// A raw stream of one of the types handled by this extension.
    #[derive(Copy, Clone, PartialEq, Eq)]
    pub enum NetworkStreamType {
      $(
        $( #[ $meta ] )?
        $i,
      )*
    }

    /// A raw stream listener of one of the types handled by this extension.
    pub enum NetworkStreamListener {
      $(
        $( #[ $meta ] )?
        $i( $listener ),
      )*
    }

    #[pin_project::pin_project(project = NetworkStreamReadHalfProject)]
    pub enum NetworkStreamReadHalf {
      $(
        $( #[ $meta ] )?
        $i( #[pin] $read_half ),
      )*
    }

    #[pin_project::pin_project(project = NetworkStreamWriteHalfProject)]
    pub enum NetworkStreamWriteHalf {
      $(
        $( #[ $meta ] )?
        $i( #[pin] $write_half ),
      )*
    }

    $(
      $( #[ $meta ] )?
      impl NetworkStreamListenerTrait for $listener {
        type Stream = $stream;
        type Addr = $addr;
        type ResourceData = ();
        const RESOURCE_NAME: &'static str = concat!(stringify!($il), "Listener");
        async fn accept(&self) -> std::io::Result<(Self::Stream, Self::Addr)> {
          <$listener> :: accept(self).await
        }
        fn listen_address(&self) -> std::io::Result<Self::Addr> {
          self.local_addr()
        }
      }

      $( #[ $meta ] )?
      impl From<$listener> for NetworkStreamListener {
        fn from(value: $listener) -> Self {
          Self::$i(value)
        }
      }

      $( #[ $meta ] )?
      impl NetworkStreamTrait for $stream {
        type Resource = $stream_resource;
        const RESOURCE_NAME: &'static str = concat!(stringify!($il), "Stream");
        fn local_address(&self) -> Result<NetworkStreamAddress, std::io::Error> {
          Ok(NetworkStreamAddress::from(self.local_addr()?))
        }
        fn peer_address(&self) -> Result<NetworkStreamAddress, std::io::Error> {
          Ok(NetworkStreamAddress::from(self.peer_addr()?))
        }
      }

      $( #[ $meta ] )?
      impl From<$stream> for NetworkStream {
        fn from(value: $stream) -> Self {
          Self::$i(value)
        }
      }
    )*

    impl NetworkStream {
      pub fn local_address(&self) -> Result<NetworkStreamAddress, std::io::Error> {
        match self {
          $(
            $( #[ $meta ] )?
            Self::$i(stm) => Ok(NetworkStreamAddress::from(stm.local_addr()?)),
          )*
        }
      }

      pub fn peer_address(&self) -> Result<NetworkStreamAddress, std::io::Error> {
        match self {
          $(
            $( #[ $meta ] )?
            Self::$i(stm) => Ok(NetworkStreamAddress::from(stm.peer_addr()?)),
          )*
        }
      }

      pub fn stream(&self) -> NetworkStreamType {
        match self {
          $(
            $( #[ $meta ] )?
            Self::$i(_) => NetworkStreamType::$i,
          )*
        }
      }

      pub fn into_split(self) -> (NetworkStreamReadHalf, NetworkStreamWriteHalf) {
        match self {
          $(
            $( #[ $meta ] )?
            Self::$i(stm) => {
              let (r, w) = stm.into_split();
              (
                NetworkStreamReadHalf::$i(r),
                NetworkStreamWriteHalf::$i(w),
              )
            },
          )*
        }
      }
    }

    impl tokio::io::AsyncRead for NetworkStream {
      fn poll_read(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
      ) -> std::task::Poll<std::io::Result<()>> {
        match self.project() {
          $(
            $( #[ $meta ] )?
            NetworkStreamProject::$i(s) => s.poll_read(cx, buf),
          )*
        }
      }
    }

    impl tokio::io::AsyncWrite for NetworkStream {
      fn poll_write(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
      ) -> std::task::Poll<Result<usize, std::io::Error>> {
        match self.project() {
          $(
            $( #[ $meta ] )?
            NetworkStreamProject::$i(s) => s.poll_write(cx, buf),
          )*
        }
      }

      fn poll_flush(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
      ) -> std::task::Poll<Result<(), std::io::Error>> {
        match self.project() {
          $(
            $( #[ $meta ] )?
            NetworkStreamProject::$i(s) => s.poll_flush(cx),
          )*
        }
      }

      fn poll_shutdown(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
      ) -> std::task::Poll<Result<(), std::io::Error>> {
        match self.project() {
          $(
            $( #[ $meta ] )?
            NetworkStreamProject::$i(s) => s.poll_shutdown(cx),
          )*
        }
      }

      fn is_write_vectored(&self) -> bool {
        match self {
          $(
            $( #[ $meta ] )?
            NetworkStream::$i(s) => s.is_write_vectored(),
          )*
        }
      }

      fn poll_write_vectored(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        bufs: &[std::io::IoSlice<'_>],
      ) -> std::task::Poll<Result<usize, std::io::Error>> {
        match self.project() {
          $(
            $( #[ $meta ] )?
            NetworkStreamProject::$i(s) => s.poll_write_vectored(cx, bufs),
          )*
        }
      }
    }

    impl tokio::io::AsyncRead for NetworkStreamReadHalf {
      fn poll_read(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
      ) -> std::task::Poll<std::io::Result<()>> {
        match self.project() {
          $(
            $( #[ $meta ] )?
            NetworkStreamReadHalfProject::$i(s) => s.poll_read(cx, buf),
          )*
        }
      }
    }

    impl tokio::io::AsyncWrite for NetworkStreamWriteHalf {
      fn poll_write(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
      ) -> std::task::Poll<Result<usize, std::io::Error>> {
        match self.project() {
          $(
            $( #[ $meta ] )?
            NetworkStreamWriteHalfProject::$i(s) => s.poll_write(cx, buf),
          )*
        }
      }

      fn poll_flush(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
      ) -> std::task::Poll<Result<(), std::io::Error>> {
        match self.project() {
          $(
            $( #[ $meta ] )?
            NetworkStreamWriteHalfProject::$i(s) => s.poll_flush(cx),
          )*
        }
      }

      fn poll_shutdown(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
      ) -> std::task::Poll<Result<(), std::io::Error>> {
        match self.project() {
          $(
            $( #[ $meta ] )?
            NetworkStreamWriteHalfProject::$i(s) => s.poll_shutdown(cx),
          )*
        }
      }

      fn is_write_vectored(&self) -> bool {
        match self {
          $(
            $( #[ $meta ] )?
            NetworkStreamWriteHalf::$i(s) => s.is_write_vectored(),
          )*
        }
      }

      fn poll_write_vectored(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        bufs: &[std::io::IoSlice<'_>],
      ) -> std::task::Poll<Result<usize, std::io::Error>> {
        match self.project() {
          $(
            $( #[ $meta ] )?
            NetworkStreamWriteHalfProject::$i(s) => s.poll_write_vectored(cx, bufs),
          )*
        }
      }
    }

    impl NetworkStreamListener {
      /// Accepts a connection on this listener.
      pub async fn accept(&self) -> Result<(NetworkStream, NetworkStreamAddress), std::io::Error> {
        Ok(match self {
          $(
            $( #[ $meta ] )?
            Self::$i(s) => {
              let (stm, addr) = s.accept().await?;
              (NetworkStream::$i(stm), addr.into())
            }
          )*
        })
      }

      pub fn listen_address(&self) -> Result<NetworkStreamAddress, std::io::Error> {
        match self {
          $(
            $( #[ $meta ] )?
            Self::$i(s) => { Ok(NetworkStreamAddress::from(s.listen_address()?)) }
          )*
        }
      }

      pub fn stream(&self) -> NetworkStreamType {
        match self {
          $(
            $( #[ $meta ] )?
            Self::$i(_) => { NetworkStreamType::$i }
          )*
        }
      }

      /// Return a `NetworkStreamListener` if a resource exists for this `ResourceId` and it is currently
      /// not locked.
      pub fn take_resource(resource_table: &mut ResourceTable, listener_rid: ResourceId) -> Result<NetworkStreamListener, JsErrorBox> {
        $(
          $( #[ $meta ] )?
          if let Some(resource) = NetworkListenerResource::<$listener>::take(resource_table, listener_rid)? {
            return Ok(resource)
          }
        )*
        Err(JsErrorBox::from_err(ResourceError::BadResourceId))
      }
    }
  };
}

network_stream!(
  [
    Tcp,
    tcp,
    tokio::net::TcpStream,
    crate::tcp::TcpListener,
    std::net::SocketAddr,
    TcpStreamResource,
    tokio::net::tcp::OwnedReadHalf,
    tokio::net::tcp::OwnedWriteHalf,
  ],
  [
    Tls,
    tls,
    crate::ops_tls::TlsStream<tokio::net::TcpStream>,
    crate::ops_tls::TlsListener,
    std::net::SocketAddr,
    TlsStreamResource,
    crate::ops_tls::TlsStreamRead<tokio::net::TcpStream>,
    crate::ops_tls::TlsStreamWrite<tokio::net::TcpStream>,
  ],
  #[cfg(unix)]
  [
    Unix,
    unix,
    tokio::net::UnixStream,
    tokio::net::UnixListener,
    tokio::net::unix::SocketAddr,
    crate::io::UnixStreamResource,
    tokio::net::unix::OwnedReadHalf,
    tokio::net::unix::OwnedWriteHalf,
  ],
  #[cfg(any(target_os = "android", target_os = "linux", target_os = "macos"))]
  [
    Vsock,
    vsock,
    tokio_vsock::VsockStream,
    tokio_vsock::VsockListener,
    tokio_vsock::VsockAddr,
    crate::io::VsockStreamResource,
    tokio_vsock::OwnedReadHalf,
    tokio_vsock::OwnedWriteHalf,
  ],
  [
    Tunnel,
    tunnel,
    crate::tunnel::TunnelStream,
    crate::tunnel::TunnelConnection,
    crate::tunnel::TunnelAddr,
    crate::tunnel::TunnelStreamResource,
    crate::tunnel::OwnedReadHalf,
    crate::tunnel::OwnedWriteHalf,
  ]
);

pub enum NetworkStreamAddress {
  Ip(std::net::SocketAddr),
  #[cfg(unix)]
  Unix(tokio::net::unix::SocketAddr),
  #[cfg(any(target_os = "android", target_os = "linux", target_os = "macos"))]
  Vsock(tokio_vsock::VsockAddr),
  Tunnel(crate::tunnel::TunnelAddr),
}

impl From<std::net::SocketAddr> for NetworkStreamAddress {
  fn from(value: std::net::SocketAddr) -> Self {
    NetworkStreamAddress::Ip(value)
  }
}

#[cfg(unix)]
impl From<tokio::net::unix::SocketAddr> for NetworkStreamAddress {
  fn from(value: tokio::net::unix::SocketAddr) -> Self {
    NetworkStreamAddress::Unix(value)
  }
}

#[cfg(any(target_os = "android", target_os = "linux", target_os = "macos"))]
impl From<tokio_vsock::VsockAddr> for NetworkStreamAddress {
  fn from(value: tokio_vsock::VsockAddr) -> Self {
    NetworkStreamAddress::Vsock(value)
  }
}

impl From<crate::tunnel::TunnelAddr> for NetworkStreamAddress {
  fn from(value: crate::tunnel::TunnelAddr) -> Self {
    NetworkStreamAddress::Tunnel(value)
  }
}

#[derive(Debug, thiserror::Error, deno_error::JsError)]
pub enum TakeNetworkStreamError {
  #[class("Busy")]
  #[error("TCP stream is currently in use")]
  TcpBusy,
  #[class("Busy")]
  #[error("TLS stream is currently in use")]
  TlsBusy,
  #[cfg(unix)]
  #[class("Busy")]
  #[error("Unix socket is currently in use")]
  UnixBusy,
  #[cfg(any(target_os = "android", target_os = "linux", target_os = "macos"))]
  #[class("Busy")]
  #[error("Vsock socket is currently in use")]
  VsockBusy,
  #[class("Busy")]
  #[error("Tunnel socket is currently in use")]
  TunnelBusy,
  #[class(generic)]
  #[error(transparent)]
  ReuniteTcp(#[from] tokio::net::tcp::ReuniteError),
  #[cfg(unix)]
  #[class(generic)]
  #[error(transparent)]
  ReuniteUnix(#[from] tokio::net::unix::ReuniteError),
  #[cfg(any(target_os = "android", target_os = "linux", target_os = "macos"))]
  #[class(generic)]
  #[error("Cannot reunite halves from different streams")]
  ReuniteVsock,
  #[class(inherit)]
  #[error(transparent)]
  Resource(deno_core::error::ResourceError),
}

/// In some cases it may be more efficient to extract the resource from the resource table and use it directly (for example, an HTTP server).
/// This method will extract a stream from the resource table and return it, unwrapped.
pub fn take_network_stream_resource(
  resource_table: &mut ResourceTable,
  stream_rid: ResourceId,
) -> Result<NetworkStream, TakeNetworkStreamError> {
  // The stream we're attempting to unwrap may be in use somewhere else. If that's the case, we cannot proceed
  // with the process of unwrapping this connection, so we just return a bad resource error.
  // See also: https://github.com/denoland/deno/pull/16242

  if let Ok(resource_rc) = resource_table.take::<TcpStreamResource>(stream_rid)
  {
    // This TCP connection might be used somewhere else.
    let resource = Rc::try_unwrap(resource_rc)
      .map_err(|_| TakeNetworkStreamError::TcpBusy)?;
    let (read_half, write_half) = resource.into_inner();
    let tcp_stream = read_half.reunite(write_half)?;
    return Ok(NetworkStream::Tcp(tcp_stream));
  }

  if let Ok(resource_rc) = resource_table.take::<TlsStreamResource>(stream_rid)
  {
    // This TLS connection might be used somewhere else.
    let resource = Rc::try_unwrap(resource_rc)
      .map_err(|_| TakeNetworkStreamError::TlsBusy)?;
    let tls_stream = resource.into_tls_stream();
    return Ok(NetworkStream::Tls(tls_stream));
  }

  #[cfg(unix)]
  if let Ok(resource_rc) =
    resource_table.take::<crate::io::UnixStreamResource>(stream_rid)
  {
    // This UNIX socket might be used somewhere else.
    let resource = Rc::try_unwrap(resource_rc)
      .map_err(|_| TakeNetworkStreamError::UnixBusy)?;
    let (read_half, write_half) = resource.into_inner();
    let unix_stream = read_half.reunite(write_half)?;
    return Ok(NetworkStream::Unix(unix_stream));
  }

  #[cfg(any(target_os = "android", target_os = "linux", target_os = "macos"))]
  if let Ok(resource_rc) =
    resource_table.take::<crate::io::VsockStreamResource>(stream_rid)
  {
    // This Vsock socket might be used somewhere else.
    let resource = Rc::try_unwrap(resource_rc)
      .map_err(|_| TakeNetworkStreamError::VsockBusy)?;
    let (read_half, write_half) = resource.into_inner();
    if !read_half.is_pair_of(&write_half) {
      return Err(TakeNetworkStreamError::ReuniteVsock);
    }
    let vsock_stream = read_half.unsplit(write_half);
    return Ok(NetworkStream::Vsock(vsock_stream));
  }

  if let Ok(resource_rc) =
    resource_table.take::<crate::tunnel::TunnelStreamResource>(stream_rid)
  {
    let resource = Rc::try_unwrap(resource_rc)
      .map_err(|_| TakeNetworkStreamError::TunnelBusy)?;
    let stream = resource.into_inner();
    return Ok(NetworkStream::Tunnel(stream));
  }

  Err(TakeNetworkStreamError::Resource(
    ResourceError::BadResourceId,
  ))
}

/// In some cases it may be more efficient to extract the resource from the resource table and use it directly (for example, an HTTP server).
/// This method will extract a stream from the resource table and return it, unwrapped.
pub fn take_network_stream_listener_resource(
  resource_table: &mut ResourceTable,
  listener_rid: ResourceId,
) -> Result<NetworkStreamListener, JsErrorBox> {
  NetworkStreamListener::take_resource(resource_table, listener_rid)
}
